"""
EPS2到SEGY完整流程模块
整合EPS2转SAC、SAC时间同步和SAC转SEGY功能
"""

import os
import sys
import argparse
from pathlib import Path

# 导入自定义模块
import eps2_to_sac
import sac_sync
import sac_to_segy


def eps2_to_segy_workflow(eps2_files, output_dir, sync_time=True, combine_segy=True, strict_one_trace_per_sac=False, split_output_files=False):
    """
    完整的EPS2到SEGY转换流程
    
    参数:
        eps2_files (list): EPS2文件路径列表
        output_dir (str): 输出目录
        sync_time (bool): 是否进行时间同步
        combine_segy (bool): 是否合并为单个SEGY文件
        strict_one_trace_per_sac (bool): 是否启用严格模式，一个SAC文件对应一个SEGY道
        split_output_files (bool): 是否按样本窗口分割为多个SEGY文件
        
    返回:
        dict: 包含各阶段输出路径的字典
    """
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 阶段1: EPS2转SAC
    sac_dir = os.path.join(output_dir, "sac")
    print(f"阶段1: 将EPS2文件转换为SAC格式...")
    sac_files = []
    
    for eps2_file in eps2_files:
        # 为每个EPS2文件创建单独的输出目录
        eps2_name = Path(eps2_file).stem
        eps2_sac_dir = os.path.join(sac_dir, eps2_name)
        os.makedirs(eps2_sac_dir, exist_ok=True)
        
        # 转换EPS2到SAC
        try:
            converted_files = eps2_to_sac.convert_eps2_to_sac(eps2_file, eps2_sac_dir)
            sac_files.extend(converted_files)
            print(f"  已转换 {eps2_file} -> {len(converted_files)} 个SAC文件")
        except Exception as e:
            print(f"  错误: 转换 {eps2_file} 失败: {e}")
            continue
    
    if not sac_files:
        raise ValueError("没有成功转换任何EPS2文件")
    
    print(f"阶段1完成: 共生成 {len(sac_files)} 个SAC文件")
    
    # 阶段2: SAC时间同步
    if sync_time:
        sync_dir = os.path.join(output_dir, "sync")
        print(f"\n阶段2: 同步SAC文件时间...")
        
        # 按分量组同步
        # 需要递归搜索子目录中的SAC文件
        synced_files = []
        for subdir in os.listdir(sac_dir):
            subdir_path = os.path.join(sac_dir, subdir)
            if os.path.isdir(subdir_path):
                # 对每个子目录进行同步
                subdir_sync_dir = os.path.join(sync_dir, subdir)
                os.makedirs(subdir_sync_dir, exist_ok=True)
                
                # 获取子目录中的SAC文件
                subdir_sac_files = [os.path.join(subdir_path, f) for f in os.listdir(subdir_path) if f.endswith('.sac')]
                
                if len(subdir_sac_files) >= 3:  # 至少需要3个文件（X、Y、Z）
                    try:
                        subdir_synced_files = sac_sync.synchronize_sac_files(subdir_sac_files, subdir_sync_dir)
                        synced_files.extend(subdir_synced_files)
                        print(f"  已同步分量组 {subdir}: {len(subdir_synced_files)} 个文件")
                    except Exception as e:
                        print(f"  错误: 同步分量组 {subdir} 失败: {e}")
        
        if not synced_files:
            raise ValueError("没有成功同步任何SAC文件")
            
        print(f"阶段2完成: 共同步 {len(synced_files)} 个SAC文件")
        
        # 使用同步后的文件进行下一步
        sac_files_for_segy = synced_files
        sac_dir_for_segy = sync_dir
    else:
        # 不进行时间同步，直接使用转换后的SAC文件
        sac_files_for_segy = sac_files
        sac_dir_for_segy = sac_dir
        print("\n阶段2: 跳过时间同步")
    
    # 阶段3: SAC转SEGY
    segy_dir = os.path.join(output_dir, "segy")
    os.makedirs(segy_dir, exist_ok=True)
    print(f"\n阶段3: 将SAC文件转换为SEGY格式...")
    
    segy_files = []
    
    if combine_segy:
        # 合并为单个SEGY文件
        # 按分量组处理
        component_groups = {}
        
        # 如果进行了时间同步，文件结构已经重新组织
        if sync_time:
            # 同步后的文件结构：sync_dir/分量组名/sync_文件名.sac
            # 遍历同步目录的子目录
            for subdir in os.listdir(sac_dir_for_segy):
                subdir_path = os.path.join(sac_dir_for_segy, subdir)
                if os.path.isdir(subdir_path):
                    # 查找该目录中的SAC文件
                    sac_files_in_subdir = [os.path.join(subdir_path, f) for f in os.listdir(subdir_path) if f.endswith('.sac')]
                    
                    # 按分量分组
                    for sac_file in sac_files_in_subdir:
                        base_name = os.path.splitext(os.path.basename(sac_file))[0]
                        # 提取文件名中的分量标识
                        if base_name.endswith('_X') or base_name.endswith('_Y') or base_name.endswith('_Z'):
                            # 对于同步后的文件，文件名格式为 sync_原文件名_X
                            # 需要提取原文件名部分
                            if base_name.startswith('sync_'):
                                common_part = base_name[5:-2]  # 去除"sync_"前缀和"_X"后缀
                            else:
                                common_part = base_name[:-2]  # 直接去除"_X"后缀
                            component = base_name[-1]  # 获取分量标识
                            
                            if common_part not in component_groups:
                                component_groups[common_part] = {}
                            component_groups[common_part][component] = sac_file
        else:
            # 没有进行时间同步，使用原始文件结构
            for sac_file in sac_files_for_segy:
                base_name = os.path.splitext(os.path.basename(sac_file))[0]
                # 提取文件名中的公共部分（去除分量标识）
                if base_name.endswith('_X') or base_name.endswith('_Y') or base_name.endswith('_Z'):
                    common_part = base_name[:-2]  # 去除最后的"_X"、"_Y"或"_Z"
                    component = base_name[-1]      # 获取分量标识
                    
                    if common_part not in component_groups:
                        component_groups[common_part] = {}
                    component_groups[common_part][component] = sac_file
        
        # 为每个分量组创建SEGY文件
        for common_part, components in component_groups.items():
            # 确保有X、Y、Z三个分量
            if not all(c in components for c in ['X', 'Y', 'Z']):
                print(f"  警告: 分量组 {common_part} 缺少完整的X、Y、Z分量，跳过")
                continue
            
            # 创建该分量组的SAC文件目录
            group_sac_dir = os.path.join(sac_dir_for_segy, common_part)
            segy_file = os.path.join(segy_dir, f"{common_part}.segy")
            
            try:
                # 转换SAC到SEGY
                sac_files_list = [Path(f) for f in [components['X'], components['Y'], components['Z']]]
                segy_path = Path(segy_file)
                # 确保输出目录存在
                segy_path.parent.mkdir(parents=True, exist_ok=True)
                sac_to_segy.combine_sacs_to_single_segy(
                    sac_files_list, segy_path, 
                    strict_one_trace_per_sac=strict_one_trace_per_sac,
                    split_output_files=split_output_files
                )
                
                # 收集生成的SEGY文件
                if split_output_files:
                    # 当split_output_files为True时，会生成多个文件
                    # 查找所有生成的文件
                    base_pattern = f"{segy_path.stem}_part*"
                    for part_file in segy_path.parent.glob(base_pattern):
                        segy_files.append(str(part_file))
                        print(f"  已创建SEGY文件: {part_file}")
                else:
                    # 单个文件
                    segy_files.append(segy_file)
                    print(f"  已创建SEGY文件: {segy_file}")
            except Exception as e:
                print(f"  错误: 创建SEGY文件 {segy_file} 失败: {e}")
    else:
        # 不合并，每个SAC文件转换为单独的SEGY文件
        for sac_file in sac_files_for_segy:
            base_name = os.path.splitext(os.path.basename(sac_file))[0]
            segy_file = os.path.join(segy_dir, f"{base_name}.segy")
            
            try:
                # 转换单个SAC到SEGY
                sac_to_segy.convert_single_sac_to_output(sac_file, segy_file)
                segy_files.append(segy_file)
                print(f"  已创建SEGY文件: {segy_file}")
            except Exception as e:
                print(f"  错误: 创建SEGY文件 {segy_file} 失败: {e}")
    
    if not segy_files:
        raise ValueError("没有成功创建任何SEGY文件")
    
    print(f"阶段3完成: 共生成 {len(segy_files)} 个SEGY文件")
    
    # 返回各阶段输出路径
    return {
        "sac_dir": sac_dir,
        "sync_dir": sync_dir if sync_time else None,
        "segy_dir": segy_dir,
        "sac_files": sac_files,
        "synced_files": synced_files if sync_time else None,
        "segy_files": segy_files
    }


def main():
    parser = argparse.ArgumentParser(description="EPS2到SEGY完整转换流程")
    parser.add_argument("input", nargs="+", help="输入EPS2文件路径或目录")
    parser.add_argument("output", help="输出目录")
    parser.add_argument("--no-sync", action="store_true", help="跳过时间同步")
    parser.add_argument("--no-combine", action="store_true", help="不合并为单个SEGY文件")
    parser.add_argument("--strict", action="store_true", help="启用严格模式，一个SAC文件对应一个SEGY道")
    parser.add_argument("--split-output-files", action="store_true", help="按样本窗口分割为多个SEGY文件")
    parser.add_argument("--recursive", action="store_true", help="递归搜索EPS2文件")
    
    args = parser.parse_args()
    
    # 收集所有EPS2文件
    eps2_files = []
    
    for input_path in args.input:
        if os.path.isfile(input_path):
            # 输入是文件
            if input_path.lower().endswith('.012'):
                eps2_files.append(input_path)
            else:
                print(f"警告: 跳过非EPS2文件: {input_path}")
        elif os.path.isdir(input_path):
            # 输入是目录，搜索EPS2文件
            if args.recursive:
                import glob
                pattern = os.path.join(input_path, "**", "*.012")
                found_files = glob.glob(pattern, recursive=True)
            else:
                import glob
                pattern = os.path.join(input_path, "*.012")
                found_files = glob.glob(pattern)
            
            if found_files:
                eps2_files.extend(found_files)
                print(f"在目录 {input_path} 中找到 {len(found_files)} 个EPS2文件")
            else:
                print(f"警告: 在目录 {input_path} 中没有找到EPS2文件")
        else:
            print(f"错误: 路径不存在: {input_path}")
    
    if not eps2_files:
        print("错误: 没有找到任何EPS2文件")
        return 1
    
    print(f"共找到 {len(eps2_files)} 个EPS2文件")
    
    # 执行完整流程
    try:
        result = eps2_to_segy_workflow(
            eps2_files, 
            args.output,
            sync_time=not args.no_sync,
            combine_segy=not args.no_combine,
            strict_one_trace_per_sac=args.strict,
            split_output_files=getattr(args, 'split_output_files', False)
        )
        
        print("\n转换流程完成!")
        print(f"SAC文件输出目录: {result['sac_dir']}")
        if result['sync_dir']:
            print(f"同步后SAC文件输出目录: {result['sync_dir']}")
        print(f"SEGY文件输出目录: {result['segy_dir']}")
        print(f"生成的SEGY文件:")
        for segy_file in result['segy_files']:
            print(f"  {segy_file}")
        
        return 0
    except Exception as e:
        print(f"错误: 转换流程失败: {e}")
        return 1


if __name__ == "__main__":
    sys.exit(main())